package com.peng.tree.test_code;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 多叉树转换为二叉树的表现形式
 * 本题测试链接：https://leetcode.com/problems/encode-n-ary-tree-to-binary-tree
 */
public class Test01_EncodeNaryTreeToBinaryTree {
    public static void main(String[] args) {
        Node node1 = new Node(1);
        Node node2 = new Node(2);
        Node node3 = new Node(3);
        Node node4 = new Node(4);

        Node node5 = new Node(5);
        Node node6 = new Node(6);
        Node node7 = new Node(7);
        Node node8 = new Node(8);
        Node node9 = new Node(9);
        Node node10 = new Node(10);

        node2.children = Arrays.asList(node5, node6);
        node3.children = Arrays.asList(node7, node8);
        node4.children = Arrays.asList(node9, node10);

        node1.children = Arrays.asList(node2, node3, node4);

        /*

                   1
             2     3    4
            5 6   7 8  9 10

         */

        Codec codec = new Codec();
        TreeNode encode = codec.encode(node1);
        System.out.println(encode);

        Node decode = codec.decode(encode);
        System.out.println(decode);


    }

    ;

    public static class Node {
        public int val;
        public List<Node> children;

        public Node() {
        }

        public Node(int _val) {
            val = _val;
        }

        public Node(int _val, List<Node> _children) {
            val = _val;
            children = _children;
        }
    }

    public static class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        TreeNode(int x) {
            val = x;
        }
    }

    static class Codec {
        private static List<Node> de(TreeNode head) {
            List<Node> children = new ArrayList<>();
            TreeNode cur = head.left;
            if (cur == null) {
                return null;
            }
            while (cur != null) {
                children.add(new Node(cur.val, de(cur)));
                cur = cur.right;
            }

            return children;
        }

        // Encodes an n-ary tree to a binary tree.
        public TreeNode encode(Node root) {
            if (root == null) {
                return null;
            }
            TreeNode head = new TreeNode(root.val);
            head.left = en(root.children);
            return head;
        }

        private TreeNode en(List<Node> children) {
            if (children == null || children.isEmpty()) {
                return null;
            }
            TreeNode head = null;
            TreeNode pre = null;
            for (Node child : children) {
                TreeNode treeNode = new TreeNode(child.val);
                treeNode.left = en(child.children);
                if (head == null) {
                    head = treeNode;
                } else {
                    pre.right = treeNode;
                }

                pre = treeNode;

            }
            return head;
        }

        // Decodes your binary tree to an n-ary tree.
        public Node decode(TreeNode root) {
            if (root == null) {
                return null;
            }
            return new Node(root.val, de(root));
        }
    }
}
